package com.witmore.neutron.ai.web.tool;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.stereotype.Service;
import technology.tabula.*;
import technology.tabula.extractors.BasicExtractionAlgorithm;

import java.io.File;
import java.io.IOException;
import java.util.List;

@Service
public class PdfToMarkdownService {

    public String convertPdfToMarkdown(File pdfFile) throws IOException {
        StringBuilder markdown = new StringBuilder();
        
        try (PDDocument document = PDDocument.load(pdfFile)) {
            // 1. 提取普通文本
            PDFTextStripper stripper = new PDFTextStripper();
            String text = stripper.getText(document);
            markdown.append(convertTextToMarkdown(text));
            
            // 2. 提取表格数据
            markdown.append("\n\n").append(extractTablesFromPdf(pdfFile));
        }
        
        return markdown.toString();
    }
    
    private String convertTextToMarkdown(String text) {
        // 基本文本转换规则
        return text.replaceAll("(?m)^\\s*$", "\n")  // 空行处理
                  .replaceAll("(?m)^(.+)$", "$1  ") // 普通行尾加两个空格保持换行
                  .replaceAll("(?m)^# (.+)$", "# $1") // 标题识别
                  .replaceAll("(?m)^## (.+)$", "## $1")
                  .replaceAll("(?m)^### (.+)$", "### $1");
    }
    
    private String extractTablesFromPdf(File pdfFile) throws IOException {
        StringBuilder tableMarkdown = new StringBuilder("\n## 表格数据\n");

        // 1. 首先加载PDF文档
        try (PDDocument document = PDDocument.load(pdfFile)) {
            // 2. 使用PDDocument初始化ObjectExtractor
            ObjectExtractor oe = new ObjectExtractor(document);
            BasicExtractionAlgorithm bea = new BasicExtractionAlgorithm();

            // 正确方式：使用PageIterator遍历页面
            PageIterator pages = oe.extract();
            while (pages.hasNext()) {
                Page page = pages.next();
                List<Table> tables = bea.extract(page);

                for (Table table : tables) {
                    tableMarkdown.append(convertTableToMarkdown(table)).append("\n");
                }
            }
        }
        
        return tableMarkdown.toString();
    }
    
    private String convertTableToMarkdown(Table table) {
        StringBuilder mdTable = new StringBuilder("\n");
        
        List<List<RectangularTextContainer>> rows = table.getRows();
        if (rows.isEmpty()) return "";
        
        // 表头
        mdTable.append("|");
        for (RectangularTextContainer cell : rows.get(0)) {
            mdTable.append(" ").append(cell.getText().replace("|", "\\|")).append(" |");
        }
        mdTable.append("\n");
        
        // 分隔线
        mdTable.append("|");
        for (int i = 0; i < rows.get(0).size(); i++) {
            mdTable.append(" --- |");
        }
        mdTable.append("\n");
        
        // 表格内容
        for (int i = 1; i < rows.size(); i++) {
            mdTable.append("|");
            for (RectangularTextContainer cell : rows.get(i)) {
                mdTable.append(" ").append(cell.getText().replace("|", "\\|")).append(" |");
            }
            mdTable.append("\n");
        }
        
        return mdTable.toString();
    }
}